"""
数据源基类模块
"""
import os
import pandas as pd
from abc import ABC, abstractmethod
from loguru import logger


class Feed(ABC):
    """
    数据源基类，定义了获取数据的接口
    """
    
    # 默认必需的数据字段
    REQUIRED_FIELDS = ['datetime', 'symbol', 'open', 'close', 'high', 'low', 'volume']
    
    def __init__(self, timeline_cache_path=None):
        """
        初始化数据源
        
        Parameters
        ----------
        timeline_cache_path : str, optional
            时间轴缓存文件路径，如果文件存在则从文件加载时间轴
        """
        # 交易品种列表
        self._instruments = []
        
        # 主控时间轴
        self._timeline = None
        self._timeline_length = 0
        self._current_time_idx = 0
        
        # 底层数据存储，使用字典存储每个品种的DataFrame
        self._df = {}
        
        # 加载时间轴缓存
        if timeline_cache_path and os.path.exists(timeline_cache_path):
            try:
                _ = pd.read_pickle(timeline_cache_path)
                self.set_timeline(_)
                self._timeline_length = len(self._timeline)
                logger.info(f"从缓存加载时间轴，共{self._timeline_length}个时间点")
            except Exception as e:
                logger.error(f"加载时间轴缓存失败: {str(e)}")
                self._timeline = None
                self._timeline_length = 0
    
    def generate_timeline(self, cache_path=None):
        """
        从所有数据源中生成统一的时间轴
        
        该函数会从所有已加载的数据中提取时间点，进行合并、去重和排序，
        生成一个统一的时间轴。如果提供了cache_path参数，还会将时间轴
        保存到指定路径。
        
        Parameters
        ----------
        cache_path : str, optional
            时间轴缓存文件路径，如果提供则将生成的时间轴保存到文件
            
        Returns
        -------
        bool
            是否成功生成时间轴
        """
        if not self._df:
            logger.warning("没有可用的数据源，无法生成时间轴")
            return False
            
        try:
            # 收集所有时间点
            all_times = []
            for instrument, df in self._df.items():
                if not isinstance(df.index, pd.DatetimeIndex):
                    logger.error(f"品种 {instrument} 的数据索引不是datetime类型")
                    return False
                all_times.extend(df.index)
            
            # 转换为Series并去重排序
            timeline = pd.Series(all_times, dtype='datetime64[ns]')
            timeline = timeline.sort_values()
            timeline = timeline.drop_duplicates()
            timeline = timeline.reset_index(drop=True)

            # 设置时间轴
            self.set_timeline(timeline)
            
            # 保存时间轴缓存
            if cache_path:
                try:
                    timeline.to_pickle(cache_path)
                    logger.info(f"时间轴缓存已保存到: {cache_path}")
                except Exception as e:
                    logger.error(f"保存时间轴缓存失败: {str(e)}")
                    return False
            return True
            
        except Exception as e:
            logger.error(f"生成时间轴失败: {str(e)}")
            return False
    
    def _check_data_fields(self):
        """检查数据字段是否符合要求"""
        for instrument, df in self._df.items():
            missing_fields = [field for field in self.REQUIRED_FIELDS if field not in df.columns]
            if missing_fields:
                logger.error(f"品种 {instrument} 数据缺少必需字段: {missing_fields}")
                raise ValueError(f"数据必须包含以下字段: {self.REQUIRED_FIELDS}")
            
            # 检查字段名是否一致（大小写、空格等）
            field_issues = []
            for field in self.REQUIRED_FIELDS:
                similar_fields = [col for col in df.columns if col.lower().strip() == field.lower()]
                if similar_fields and similar_fields[0] != field:
                    field_issues.append(f"{similar_fields[0]} 应该改为 {field}")
            
            if field_issues:
                logger.warning(f"品种 {instrument} 数据字段名称不标准: {field_issues}")
    
    @property
    def instruments(self):
        """获取当前数据源包含的所有交易品种列表"""
        return self._instruments
    
    @property
    def timeline(self):
        """获取时间轴"""
        return self._timeline
    
    @property
    def current_time(self):
        """获取当前时间"""
        if self._timeline is None:
            return None
        try:
            return self._timeline.iloc[self._current_time_idx]
        except IndexError:
            logger.error(f"时间轴索引越界: {self._current_time_idx}")
            return None
    
    def get_current_bar(self, instrument):
        """
        获取指定品种当前时间点的数据
        
        Parameters
        ----------
        instrument : str
            交易品种代码
            
        Returns
        -------
        pandas.Series
            当前时间点的数据，如果没有数据返回None
        """
        if instrument not in self._df:
            logger.error(f"未找到品种: {instrument}")
            return None
            
        current_time = self.current_time
        if current_time is None:
            return None
            
        try:
            return self._df[instrument].loc[current_time]
        except KeyError:
            return None
    
    def set_timeline(self, timeline, cache_path=None):
        """
        设置主控时间轴
        
        Parameters
        ----------
        timeline : pandas.Series
            时间序列，包含所有可能的交易时间点
        """
        if timeline is None:
            logger.warning("设置了空的时间轴")
            return
        
        if not isinstance(timeline, pd.Series):
            logger.error("时间轴必须是pandas.Series类型")
            raise TypeError("时间轴必须是pandas.Series类型")
        
        self._timeline = timeline
        self._timeline = self._timeline.sort_values()
        self._timeline = self._timeline.drop_duplicates()
        self._timeline = self._timeline.reset_index(drop=True)
        self._timeline_length = len(timeline)
        self._current_time_idx = 0
        logger.info(f"设置时间轴，共{self._timeline_length}个时间点")
    
    def next(self):
        """
        移动到下一个时间点
        
        Returns
        -------
        bool
            如果还有下一个时间点返回True，否则返回False
        """
        if self._timeline is None:
            logger.warning("时间轴未设置，无法前进到下一个时间点")
            return False
            
        self._current_time_idx += 1
        if self._current_time_idx >= self._timeline_length:
            logger.info("已到达时间轴末尾")
            return False
        
        return True
    
    def reset(self):
        """重置数据源到初始状态"""
        self._current_time_idx = 0
        logger.info("重置数据源到初始状态")
    
    @abstractmethod
    def get_data_series(self, instrument, column_name='close'):
        """
        获取指定品种的特定数据序列
        
        Parameters
        ----------
        instrument : str
            交易品种代码
        column_name : str, optional
            数据列名，默认为'close'
            
        Returns
        -------
        pandas.Series
            数据序列
        """
        pass
